import numpy as np
from ultralytics import YOLO
import cv2
import time
from datetime import datetime
from sqlalchemy import create_engine, Table, Column, Integer, DateTime, String, Enum, MetaData, LargeBinary, insert
from sqlalchemy.orm import scoped_session, sessionmaker
import uuid

# 创建 MySQL 数据库连接
engine = create_engine('mysql+pymysql://root:root@localhost:3306/flask_sql', echo=True)
# 定义数据表
metadata = MetaData()
dan = Table('dan', metadata,
            Column('id', Integer, primary_key=True),
            Column('time', DateTime, nullable=False),
            Column('location', String(64), nullable=False),
            Column('behavior', String(64), nullable=False),
            Column('image_data', LargeBinary, nullable=False),
            Column('status', Enum('pending', 'approved', 'rejected'), nullable=False)
            )

# 创建数据表
metadata.create_all(engine)

Session = scoped_session(sessionmaker(bind=engine))

# 记录人员进入区域的时间
enter_time = None

# 配置安全距离和停留时间阈值（可根据需求调整）
SAFE_DISTANCE = 5  # 安全距离（像素）
ALERT_DURATION = 0.1  # 停留时间（秒）
# 指定区域 (例: 一个矩形区域)
#region = (100, 100, 400, 400)  # x1, y1, x2, y2
# 定义 init_model 函数，用于根据传入的模型名称初始化 YOLO 对象
def init_model(model_path):
    model = YOLO(model_path)  # TODO:  实例化 YOLO 对象，传入模型名称
    return model  # 返回 YOLO 对象实例

def is_touching_box(box, region):
    x1, y1, x2, y2 = box
    rx1, ry1, rx2, ry2 = region
    # 检查盒子是否触及区域边界
    return not (x2 < rx1 or x1 > rx2 or y2 < ry1 or y1 > ry2)

def calculate_distance_to_region(box, region):
    x1, y1, x2, y2 = box
    rx1, ry1, rx2, ry2 = region
    center_x = (x1 + x2) / 2
    center_y = (y1 + y2) / 2

    # 距离区域边框的最小距离
    distances = [
        abs(center_x - rx1),
        abs(center_x - rx2),
        abs(center_y - ry1),
        abs(center_y - ry2)
    ]
    return min(distances)

def process_frame(model1, model2, frame, region, show_box=True, show_mask=False):
    global enter_time  # 引用全局变量

    # 使用第一个 YOLO 模型对输入的视频帧进行目标检测
    results1 = model1.predict(frame, conf=0.25, iou=0.8)

    # 使用第二个 YOLO 模型对输入的视频帧进行目标检测
    results2 = model2.predict(frame, conf=0.25, iou=0.8)

    # 创建Pframe来存储处理后的帧
    Pframe = frame.copy()

    # 绘制指定区域
    cv2.rectangle(Pframe, (region[0], region[1]), (region[2], region[3]), (255,0, 0), 2)
    cv2.putText(Pframe, "bank vault", (region[0], region[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

    # 合并所有检测结果
    all_results = []

    # 处理第一个模型的检测结果
    for result in results1:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = box.conf[0]
            label = result.names[box.cls[0]]
            # 判断是否触及指定区域
            if is_touching_box((x1, y1, x2, y2), region):
                color = (0, 0, 255)  # 红色
                label = f'{label} {conf:.2f} warning'
            else:
                color = (0, 255, 0)  # 绿色
                label = f'{label} {conf:.2f}'
            all_results.append((x1, y1, x2, y2, label, color))

    # 处理第二个模型的检测结果
    for result in results2:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = box.conf[0]
            label = result.names[box.cls[0]]
            # 判断是否触及指定区域
            if is_touching_box((x1, y1, x2, y2), region):
                color = (0, 0, 255)  # 红色
                label = f'{label} {conf:.2f} warning'
            else:
                color = (0, 255, 0)  # 绿色
                label = f'{label} {conf:.2f}'
            all_results.append((x1, y1, x2, y2, label, color))

    # 绘制所有检测结果
    for x1, y1, x2, y2, label, color in all_results:
        if show_box:
            cv2.rectangle(Pframe, (x1, y1), (x2, y2), color, 2)
        cv2.putText(Pframe, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)

    return Pframe


def process_frame(model, frame, region, show_box=True, show_mask=False):
    global enter_time  # 引用全局变量

    # 使用传入的 YOLO 模型对输入的视频帧进行目标检测，设置目标检测置信度阈值为 0.25
    results = model.predict(frame, conf=0.25, iou=0.8)

    # 创建Pframe来存储处理后的帧
    Pframe = frame.copy()

    # 绘制指定区域
    if isinstance(Pframe, np.ndarray):
        if isinstance(region, (list, tuple)) and len(region) == 4:
            cv2.rectangle(Pframe, (region[0], region[1]), (region[2], region[3]), (0, 0, 255), 2)
            cv2.putText(Pframe, "bank vault", (region[0], region[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0),
                        2)
        else:
            print("Invalid region format.")
    else:
        print("Invalid frame format.")


    # 处理检测结果
    for result in results:
        boxes = result.boxes
        for box in boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = float(box.conf[0])  # 置信度
            cls = int(box.cls[0])  # 类别
            label = f"{model.names[cls]}: {conf:.2f}"

            # 判断是否进入防护区域
            if is_touching_box((x1, y1, x2, y2), region):
                color = ( 0, 0,255)  # 红色
                label = f'warning-{label}'

                # 检查人员是否进入区域
                if enter_time is None:
                    enter_time = time.time()
                    # 触发进入区域的告警，插入数据库
                    save_alert_to_db(frame, '闯入', '金库', 'pending')

                    # 计算距离区域边框的距离
                    distance_to_region = calculate_distance_to_region((x1, y1, x2, y2), region)
                    if distance_to_region < SAFE_DISTANCE:
                        color = ( 0, 0, 255)  # 红色
                        save_alert_to_db(frame, '接近', '金库', 'pending')
                    else:
                        color = (0, 255, 0)  # 绿色

                        # 检查人员在区域内的停留时间
                        if time.time() - enter_time > ALERT_DURATION:
                            # 触发停留时间超过阈值的告警
                            save_alert_to_db(frame, '停留超时', '金库', 'pending')
            else:
                # 如果人员离开了区域，重置进入时间
                enter_time = None

                # 增加判断：如果类别是 smoking，设置颜色为红色；
                # 后面增加告警机制
                if model.names[cls] == "Smoking":
                    color = ( 0, 0,255)  # 红色
                    save_alert_to_db(frame, '吸烟', '金库', 'pending')
                elif model.names[cls] == "Smoke" or model.names[cls] == "Fire":
                    color = (0, 0, 255)  # 红色
                    save_alert_to_db(frame, '明火烟雾', '金库', 'pending')
                else:
                    color = (0, 255, 0)  # 绿色

            if show_box:
                cv2.rectangle(Pframe, (x1, y1), (x2, y2), color, 2)
            cv2.putText(Pframe, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
           # 调用准备和推送数据的函数
   # prepare_and_push_data(Pframe, results)

    return Pframe

def save_alert_to_db(frame, behavior, location, status):
    session = Session()
    try:
        _, image_binary = cv2.imencode('.jpg', frame)
        image_data = image_binary.tobytes()
        current_time = datetime.now()

        stmt = insert(dan).values(
            time=current_time,
            location=location,
            behavior=behavior,
            image_data=image_data,
            status=status
        )
        session.execute(stmt)
        session.commit()
        print(f"告警数据插入成功：{behavior} at {current_time}")

    except Exception as e:
        session.rollback()
        print(f"插入告警数据失败: {e}")
    finally:
        session.close()
